#!/usr/bin/env python
#-*- coding: utf-8 -*-
#
# Copyright 2013 Antoine Drouin (poinix@gmail.com)
#
# This file is part of PAT.
#
#    PAT is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    PAT is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with PAT.  If not, see <http://www.gnu.org/licenses/>.
#

import math
import numpy as np
import scipy
import scipy.integrate
import scipy.signal
import matplotlib.pyplot as plt
import control.matlab

import pat.vehicles.rotorcraft.dynamic_model_synth as dm
import pat.vehicles.rotorcraft.autopilot_synth as sap
import pat.autopilot as ap
import pat.utils as pu


class TestLaw(ap.LinearFeedback):
    def __init__(self, dm, Xe, Ue):
      K = np.array([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
                    [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
                    [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
                    [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
                    [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])
      ap.LinearFeedback.__init__(self, dm.sv_size, dm.iv_size, Xe, Ue, K, "Test")
      

class TestLawLQR(ap.LQRFeedback):
    def __init__(self, dm, Xe, Ue, refs=()):
        A, B = dm.get_jacobian(Xe, Ue) 
        #                  x       y      z     u       v      w    phi   theta psi     p       q       r   om
        self.Q = np.diag([0.005, 0.005, 0.1, 0.005, 0.005, 0.005,  0.05, 0.05, 0.1, 0.001,  0.001,  0.001, 0.01])
        #                 lat   lon   rud   col   thr 
        self.R = np.diag([0.01, 0.01, 0.05, 0.05, 0.1])
        (K, X, E) = control.matlab.lqr(A, B, self.Q, self.R)
        ap.LQRFeedback.__init__(self, dm.sv_size, dm.iv_size, Xe, Ue, -K, "PositionLQR", refs)
        

my_dm = dm.DynamicModel()
print "Param:\n", my_dm.P

np.set_printoptions(precision=2, suppress=True, linewidth=200)
args = {'om':pu.rps_of_rpm(1700.), 'vel':2.}
Xe, Ue = my_dm.trim(args, debug=True)
print "Ue\n", np.array(Ue)
print "Xe\n", np.array(Xe)
A, B = my_dm.get_jacobian(Xe, Ue)
print "A\n", A
print "B\n", B
pu.print_lti_dynamics(A, B)
                   
time = np.arange(0., 10., 0.01)
X = np.zeros((len(time), dm.sv_size))
X0 = np.array(Xe); X0[dm.sv_x] -=5# ; X0[dm.sv_z] -=1 ; X0[dm.sv_psi] += pu.rad_of_deg(20)
X[0] = my_dm.reset(X0); 
U = Ue*np.zeros((len(time), dm.iv_size))
Xref = np.zeros((len(time), my_dm.sv_size))

my_ref = sap.RefGo2(Xe, Ue, "RefGo", period=200, dt=50)
my_law = TestLawLQR(my_dm, Xe, Ue, [my_ref])

for i in range(1,len(time)):
    U[i-1], Xref[i-1] = my_law.run(time[i-1], X[i-1], ref=my_ref)
    X[i,:] = my_dm.run(time[i]-time[i-1], U[i-1])
U[-1], Xref[-1] = my_law.run(time[-1], X[-1], ref=my_ref)
my_dm.plot_trajectory(time, X, U, Xref, window_title="heli synth 3")
plt.show()
